﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using Aeon.Emulator.DebugSupport;

namespace Aeon.Emulator.Decoding.Recompiler
{
    /// <summary>
    /// Recompiles x86 instructions into IL methods.
    /// </summary>
    internal sealed class Compiler
    {
        private static readonly HashSet<string> ConditionalJumps = new HashSet<string>()
        {
            "jo", "jno", "jc", "jnc", "jz", "jnz", "jcz", "jncz", "js", "jns", "jl", "jge", "jle", "jg", "jcxz", "jp", "jnp"
        };

        private static readonly HashSet<string> Loops = new HashSet<string>()
        {
            "loop", "loope", "loopne"
        };

        private static readonly HashSet<string> ExitInstructions = new HashSet<string>()
        {
            "call", "ret", "int", "int 3h", "into", "iret", "inth", "callh", "cmd"
        };

        private struct InstructionUnit
        {
            public InstructionUnit(Instruction inst, UnitType type, uint target)
            {
                this.Instruction = inst;
                this.UnitType = type;
                this.JumpTarget = target;
            }

            public Instruction Instruction;
            public UnitType UnitType;
            public uint JumpTarget;

            public override string ToString()
            {
                return this.Instruction.ToString();
            }
        }

        /// <summary>
        /// Describes an instruction as it relates to the compilation unit.
        /// </summary>
        private enum UnitType
        {
            /// <summary>
            /// The instruction will execute once and continue.
            /// </summary>
            Normal,
            /// <summary>
            /// The instruction will jump to an internal label.
            /// </summary>
            Jump,
            /// <summary>
            /// The instruction may jump to an internal label.
            /// </summary>
            ConditionalJump,
            /// <summary>
            /// The instruction is a loop instruction.
            /// </summary>
            Loop,
            /// <summary>
            /// The instruction is a string instruction.
            /// </summary>
            String,
            /// <summary>
            /// The instruction will end the current functional unit.
            /// </summary>
            Exit
        }

        #region Private Fields
        private readonly RegisterExpressions registers;
        private readonly VirtualMachine vm;
        private ILGenerator il;
        private TempStorage temp = new TempStorage();
        private Dictionary<uint, Label> jumpLabels = new Dictionary<uint, Label>();
        private Addresser16 address16;
        private Addresser32 address32;
        private Dictionary<string, StringCompilers.StringBase> stringCompilers;
        private LoopCompiler loopCompiler;
        private Conditions conditions;
        #endregion

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the Compiler class.
        /// </summary>
        /// <param name="vm">Associated VirtualMachine instance.</param>
        public Compiler(VirtualMachine vm)
        {
            this.vm = vm;
            this.registers = new RegisterExpressions(vm.Processor);
            this.address16 = new Addresser16(this.registers);
            this.address32 = new Addresser32(this.registers);
            this.loopCompiler = new LoopCompiler(this.registers);
            this.conditions = new Conditions(this.registers);

            var cmps = new StringCompilers.CmpsCompiler(this.registers);
            var lods = new StringCompilers.LodsCompiler(this.registers);
            var movs = new StringCompilers.MovsCompiler(this.registers);
            var outs = new StringCompilers.OutsCompiler(this.registers);
            var scas = new StringCompilers.ScasCompiler(this.registers);
            var stos = new StringCompilers.StosCompiler(this.registers);

            this.stringCompilers = new Dictionary<string, StringCompilers.StringBase>()
            {
                { "cmpsb", cmps }, { "cmpsw", cmps },
                { "lodsb", lods }, { "lodsw", lods },
                { "movsb", movs }, { "movsw", movs },
                { "outsb", outs }, { "outsw", outs },
                { "scasb", scas }, { "scasw", scas },
                { "stosb", stos }, { "stosw", stos }
            };
        }
        #endregion

        #region Public Methods
        /// <summary>
        /// Recompiles a block of instructions.
        /// </summary>
        /// <param name="segment">Segment of instructions to recompile.</param>
        /// <param name="offset">Offset of instructions to recompile.</param>
        /// <param name="maxCount">Maximum number of instructions to run.</param>
        /// <returns>Recompiled instruction block.</returns>
        public CompiledBlock Compile(ushort segment, uint offset, int maxCount)
        {
            var disasm = new Disassembler(this.vm)
            {
                StartSegment = segment,
                StartOffset = offset,
                MaximumInstructions = maxCount
            };

            var dynamicMethod = new DynamicMethod(string.Empty, typeof(int), new Type[] { typeof(VirtualMachine) }, typeof(Compiler));
            this.il = dynamicMethod.GetILGenerator();
            var context = new CompilerContext(il.DeclareLocal(typeof(int)), maxCount, this.registers);
            this.jumpLabels.Clear();
            this.temp.Initialize(this.il);

            var units = Analyze(disasm.Instructions);
            if(units.Count == 0)
                throw new InvalidOperationException();
            
            foreach(var unit in units)
            {
                Label jumpTarget;
                if(this.jumpLabels.TryGetValue(unit.Instruction.EIP, out jumpTarget))
                    il.MarkLabel(jumpTarget);

                CompileUnit(unit, context);
            }

            this.il.LoadLocal(context.CounterInt);
            this.il.Emit(OpCodes.Ret);

            var dlg = (Func<VirtualMachine, int>)dynamicMethod.CreateDelegate(typeof(Func<VirtualMachine, int>));
            var block = new CompiledBlock(dlg);
            block.SetInstructions(units.Select(u => u.Instruction));
            return block;
        }
        #endregion

        #region Private Methods
        private void CompileUnit(InstructionUnit unit, CompilerContext context)
        {
            context.SetInstruction(unit.Instruction);
            InstructionProlog(context.Instruction);
   
            switch(unit.UnitType)
            {
            case UnitType.Normal:
                CompileNormalInstruction(context);
                InstructionEpilog(context);
                break;

            case UnitType.Exit:
                CompileNormalInstruction(context);
                InstructionEpilog(context);
                il.LoadLocal(context.CounterInt);
                il.Emit(OpCodes.Ret);
                break;

            case UnitType.Loop:
                GenerateLoop(context, unit);
                break;

            case UnitType.String:
                this.stringCompilers[unit.Instruction.Opcode.Name].Generate(context, il);
                break;

            case UnitType.Jump:
                CompileJump(context, unit.JumpTarget);
                break;

            case UnitType.ConditionalJump:
                CompileConditionalJump(context, unit.JumpTarget);
                break;
            }
        }

        private void GenerateLoop(CompilerContext context, InstructionUnit unit)
        {
            Label target;
            if(this.jumpLabels.TryGetValue(unit.JumpTarget, out target))
                this.loopCompiler.GenerateLoop(unit.Instruction.Opcode.Name, context, target, unit.JumpTarget, il);
            else
                this.loopCompiler.GenerateLoop(unit.Instruction.Opcode.Name, context, null, unit.JumpTarget, il);
        }

        private void CompileConditionalJump(CompilerContext context, uint target)
        {
            var jumpLabel = il.DefineLabel();
            var continueLabel = il.DefineLabel();

            this.conditions.GenerateJump(context.Instruction.Opcode.Name, context.AddressSize, jumpLabel, il);
            il.Emit(OpCodes.Br_S, continueLabel);

            il.MarkLabel(jumpLabel);
            SetEIP(target);

            Label targetLabel;
            if(this.jumpLabels.TryGetValue(target, out targetLabel))
            {
                if(target > context.Instruction.EIP)
                    il.Emit(OpCodes.Br, targetLabel);
                else
                {
                    il.LoadLocal(context.CounterInt);
                    il.LoadConstant(context.MaxCount);
                    il.Emit(OpCodes.Clt);
                    il.Emit(OpCodes.Brtrue, targetLabel);
                    il.LoadLocal(context.CounterInt);
                    il.Emit(OpCodes.Ret);
                }
            }
            else
            {
                il.LoadLocal(context.CounterInt);
                il.Emit(OpCodes.Ret);
            }
            
            il.MarkLabel(continueLabel);
        }

        private void CompileJump(CompilerContext context, uint target)
        {
            SetEIP(target);

            Label targetLabel;
            if(this.jumpLabels.TryGetValue(target, out targetLabel))
            {
                if(target > context.Instruction.EIP)
                    il.Emit(OpCodes.Br, targetLabel);
                else
                {
                    il.LoadLocal(context.CounterInt);
                    il.LoadConstant(context.MaxCount);
                    il.Emit(OpCodes.Clt);
                    il.Emit(OpCodes.Brtrue, targetLabel);
                    il.LoadLocal(context.CounterInt);
                    il.Emit(OpCodes.Ret);
                }
            }
            else
            {
                il.LoadLocal(context.CounterInt);
                il.Emit(OpCodes.Ret);
            }
        }

        private void SetEIP(uint newEIP)
        {
            this.registers.GenerateReference("EIP", il);
            il.LoadConstant((int)newEIP);
            il.Emit(OpCodes.Stind_I4);
        }

        private List<InstructionUnit> Analyze(IEnumerable<Instruction> instructions)
        {
            var units = new List<InstructionUnit>();

            foreach(var inst in instructions)
            {
                if(inst.Opcode == null)
                    break;

                var unit = AnalyzeInstruction(inst);
                units.Add(unit);
            }

            var reachable = new HashSet<uint>();

            VisitAllInstructions(units, 0, reachable);
            units.RemoveAll(u => !reachable.Contains(u.Instruction.EIP));

            foreach(var unit in units)
            {
                if(unit.UnitType == UnitType.ConditionalJump || unit.UnitType == UnitType.Jump || unit.UnitType == UnitType.Loop)
                {
                    if(reachable.Contains(unit.JumpTarget))
                    {
                        if(!this.jumpLabels.ContainsKey(unit.JumpTarget))
                            this.jumpLabels.Add(unit.JumpTarget, il.DefineLabel());
                    }
                }
            }

            return units;
        }

        private static void VisitAllInstructions(List<InstructionUnit> units, int offset, HashSet<uint> reachable)
        {
            for(int i = offset; i < units.Count; i++)
            {
                if(reachable.Contains(units[i].Instruction.EIP))
                    return;

                reachable.Add(units[i].Instruction.EIP);

                if(units[i].UnitType == UnitType.Jump || units[i].UnitType == UnitType.ConditionalJump || units[i].UnitType == UnitType.Loop)
                {
                    int index = GetListIndex(units, units[i].JumpTarget);
                    if(index >= 0)
                        VisitAllInstructions(units, index, reachable);
                }

                if(units[i].UnitType == UnitType.Jump || units[i].UnitType == UnitType.Exit)
                    return;
            }
        }

        private static int GetListIndex(List<InstructionUnit> units, uint eip)
        {
            for(int i = 0; i < units.Count; i++)
            {
                if(units[i].Instruction.EIP == eip)
                    return i;
            }

            return -1;
        }

        private InstructionUnit AnalyzeInstruction(Instruction instruction)
        {
            var unitType = UnitType.Normal;
            uint jumpTarget = 0;

            if(this.stringCompilers.ContainsKey(instruction.Opcode.Name))
                unitType = UnitType.String;
            else if(ConditionalJumps.Contains(instruction.Opcode.Name))
            {
                unitType = UnitType.ConditionalJump;
                jumpTarget = CalculateRelativeJump(instruction.EIP + (uint)instruction.Length, instruction.Operands.Operand1.ImmediateValue, 16);
                //jumpTarget = (uint)(instruction.EIP + instruction.Length + (int)instruction.Operands.Operand1.ImmediateValue);
            }
            else if(instruction.Opcode.Name == "jmp")
            {
                unitType = UnitType.Jump;
                if(instruction.Operands.Operand1.Type == CodeOperandType.RelativeJumpAddress)
                    jumpTarget = CalculateRelativeJump(instruction.EIP + (uint)instruction.Length, instruction.Operands.Operand1.ImmediateValue, 16);
                else
                    unitType = UnitType.Exit;
            }
            else if(Loops.Contains(instruction.Opcode.Name))
            {
                unitType = UnitType.Loop;
                jumpTarget = CalculateRelativeJump(instruction.EIP + (uint)instruction.Length, (uint)(sbyte)instruction.Operands.Operand1.ImmediateValue, 16);
            }
            else if(ExitInstructions.Contains(instruction.Opcode.Name))
                unitType = UnitType.Exit;
      
            return new InstructionUnit(instruction, unitType, jumpTarget);
        }
        private uint CalculateRelativeJump(uint start, uint offset, int mode)
        {
            uint res1, res2;

            res1 = (ushort)((short)start + (short)offset);
            res2 = (uint)((int)start + (int)offset);

            //if(res1 != res2)
            //    System.Diagnostics.Debugger.Break();

            if(mode == 16)
                return res1;
            else
                return res2;
        }

        private void LoadVirtualMachine()
        {
            il.LoadArgument(0);
        }
        private void InstructionProlog(Instruction instruction)
        {
            this.registers.GenerateReference("EIP", il);
            il.LoadConstant((int)(instruction.EIP + instruction.Length));
            il.Emit(OpCodes.Stind_I4);
        }
        private void InstructionEpilog(CompilerContext context)
        {
            il.LoadLocal(context.CounterInt);
            il.LoadConstant(1);
            il.Emit(OpCodes.Add);
            il.StoreLocal(context.CounterInt);
        }
        private Action BuildMemoryOperand(CompilerContext context)
        {
            LocalBuilder tempValue;
            MethodInfo getMethod;
            MethodInfo setMethod;
            
            if(context.OperandSize == 8)
            {
                tempValue = this.temp.Byte;
                getMethod = Infos.PhysicalMemory.GetByte;
                setMethod = Infos.PhysicalMemory.SetByte;
            }
            else if(context.OperandSize == 16)
            {
                tempValue = this.temp.Short;
                getMethod = Infos.PhysicalMemory.GetUInt16;
                setMethod = Infos.PhysicalMemory.SetUInt16;
            }
            else if(context.OperandSize == 32)
            {
                tempValue = this.temp.Int;
                getMethod = Infos.PhysicalMemory.GetUInt32;
                setMethod = Infos.PhysicalMemory.SetUInt32;
            }
            else
                throw new InvalidOperationException();

            if(context.AddressSize == 16)
                this.address16.GenerateLinearAddress(context.Operand, context.Instruction.Prefixes, il);
            else
                this.address32.GenerateLinearAddress(context.Operand, context.Instruction.Prefixes, il);

            il.StoreLocal(this.temp.Address);

            if(context.FlowDirection != CodeOperandFlow.Out)
            {
                LoadVirtualMachine();
                il.Emit(OpCodes.Ldfld, Infos.VirtualMachine.PhysicalMemory);
                il.LoadLocal(this.temp.Address);
                il.Emit(OpCodes.Call, getMethod);
                il.StoreLocal(tempValue);
            }

            if(context.FlowDirection == CodeOperandFlow.In)
            {
                il.LoadLocal(tempValue);
                return null;
            }
            else
            {
                il.Emit(OpCodes.Ldloca_S, (byte)tempValue.LocalIndex);
                return () => 
                {
                    LoadVirtualMachine();
                    il.Emit(OpCodes.Ldfld, Infos.VirtualMachine.PhysicalMemory);
                    il.LoadLocal(this.temp.Address);
                    il.LoadLocal(tempValue);
                    il.Emit(OpCodes.Call, setMethod);
                };
            }
        }

        private void CompileXLAT16(CompilerContext context)
        {
            registers.GenerateReference("AL", il);

            LoadVirtualMachine();
            il.Emit(OpCodes.Ldfld, Infos.VirtualMachine.PhysicalMemory);

            var segment = Addresser16.GetOverrideSegment(context.Instruction.Prefixes) ?? "DS";
            registers.GenerateSegmentBaseValue(segment, il);

            registers.GenerateValue("BX", il);
            il.Emit(OpCodes.Conv_U4);
            
            registers.GenerateValue("AL", il);
            il.Emit(OpCodes.Conv_U4);
            il.Emit(OpCodes.Add);

            il.Emit(OpCodes.Conv_U2);
            il.Emit(OpCodes.Conv_U4);

            il.Emit(OpCodes.Add);

            il.Emit(OpCodes.Call, Infos.PhysicalMemory.GetByte);

            il.Emit(OpCodes.Stind_I1);
        }

        private void CompileNormalInstruction(CompilerContext context)
        {
            var operands = context.Instruction.Operands;
            Action writeBack = null;
            SegmentIndex? outSegment = null;

            if(context.Instruction.Opcode.Name == "xlat")
            {
                CompileXLAT16(context);
                return;
            }
            
            LoadVirtualMachine();
            
            if(operands.Operand1.Type != CodeOperandType.None)
            {
                context.SetOperand(operands.Operand1, 0);
                if(context.Instruction.Opcode.Operands[0] == OperandType.SegmentRegister && context.FlowDirection != CodeOperandFlow.In)
                    outSegment = (SegmentIndex)Enum.Parse(typeof(SegmentIndex), context.Operand.RegisterValue.ToString());
                
                writeBack = BuildOperand(context);

                if(operands.Operand2.Type != CodeOperandType.None)
                {
                    context.SetOperand(operands.Operand2, 1);
                    if(BuildOperand(context) != null)
                        throw new InvalidOperationException();

                    if(operands.Operand3.Type != CodeOperandType.None)
                    {
                        context.SetOperand(operands.Operand3, 2);
                        BuildOperand(context);
                    }
                }
            }

            int methodIndex = 0;
            if(context.OperandSize == 32)
                methodIndex |= 1;
            if(context.AddressSize == 32)
                methodIndex |= 2;

            il.Emit(OpCodes.Call, context.Instruction.Opcode.EmulateMethods[methodIndex]);

            if(writeBack != null)
                writeBack();

            if(outSegment != null)
            {
                LoadVirtualMachine();
                il.LoadConstant((int)(SegmentIndex)outSegment);
                il.Emit(OpCodes.Call, Infos.VirtualMachine.UpdateSegment);
            }
        }
        private Action BuildOperand(CompilerContext context)
        {
            var operand = context.Operand;
            switch(operand.Type)
            {
            case CodeOperandType.Immediate:
                il.LoadConstant((int)operand.ImmediateValue);
                if(context.OperandSize == 16)
                    il.Emit(OpCodes.Conv_U2);
                else if(context.OperandSize == 8)
                    il.Emit(OpCodes.Conv_U1);
                break;

            case CodeOperandType.RelativeJumpAddress:
                il.LoadConstant((int)operand.ImmediateValue);
                if(context.OperandSize == 16)
                    il.Emit(OpCodes.Conv_I2);
                break;

            case CodeOperandType.Register:
                if(context.FlowDirection == CodeOperandFlow.In)
                    this.registers.GenerateValue(operand.RegisterValue.ToString(), il);
                else
                    this.registers.GenerateReference(operand.RegisterValue.ToString(), il);
                break;

            case CodeOperandType.MemoryAddress:
            case CodeOperandType.AbsoluteJumpAddress:
                return BuildMemoryOperand(context);

            case CodeOperandType.EffectiveAddress:
                if(context.AddressSize == 16)
                    this.address16.GenerateOffset(context.Operand, il);
                else
                    this.address32.GenerateOffset(context.Operand, il);
                break;

            case CodeOperandType.FarMemoryAddress:
                if(context.AddressSize == 16)
                    il.LoadConstant((int)((operand.ImmediateValue & 0xFFFFu) | ((uint)operand.FarSegment << 16)));
                else
                    il.Emit(OpCodes.Ldc_I8, (long)(operand.ImmediateValue | ((long)operand.FarSegment << 32)));
                break;

            case CodeOperandType.IndirectFarMemoryAddress:
                LoadVirtualMachine();
                il.Emit(OpCodes.Ldfld, Infos.VirtualMachine.PhysicalMemory);
                if(context.AddressSize == 16)
                {
                    this.address16.GenerateLinearAddress(context.Operand, context.Instruction.Prefixes, il);
                    il.Emit(OpCodes.Call, Infos.PhysicalMemory.GetUInt32);
                }
                else
                {
                    this.address32.GenerateLinearAddress(context.Operand, context.Instruction.Prefixes, il);
                    il.Emit(OpCodes.Call, Infos.PhysicalMemory.GetUInt64);
                }
                break;
            }

            return null;
        }
        #endregion
    }
}
